import matplotlib.pyplot as plt
import torch
import numpy as np
# 导入必要的库
import os
import sys

# 获取当前路径的上一级目录
parent_dir = os.path.abspath(os.path.join(os.getcwd(), ".."))
print(parent_dir)
# 将上一级目录添加到 sys.path
sys.path.append(parent_dir)
# 现在可以导入 d2lzh_pytorch 模块
import d2lzh_pytorch as d2l


def semilogy_test(x_vals, y_vals, x_label, y_label, x2_vals=None, y2_vals=None, legend=None, figsize=(3.5, 2.5)):
    d2l.set_figsize(figsize)
    plt.xlabel(x_label)
    plt.ylabel(y_label)

    # 创建一个2行1列的子图
    plt.subplot(1, 1, 1)  # 1行1列，第1个子图
    plt.semilogy(x_vals, y_vals)

    # 如果传入第二组数据，则在同一个图中绘制第二条曲线
    if x2_vals is not None and y2_vals is not None:
        plt.semilogy(x2_vals, y2_vals, linestyle=':')
        plt.legend(legend)

    plt.show()


def fit_and_plot(train_features, test_features, train_labels, test_labels):
    net = torch.nn.Linear(train_features.shape[-1], 1)
    # 通过Linear⽂档可知，pytorch已经将参数初始化了，所以我们这⾥就不⼿动初始化了

    batch_size = min(10, train_labels.shape[0])
    dataset = torch.utils.data.TensorDataset(train_features, train_labels)
    #生成数据批量迭代器，并启用随机打乱功能
    train_iter = torch.utils.data.DataLoader(dataset, batch_size, shuffle=True)
    #使用随机梯度下降优化器
    optimizer = torch.optim.SGD(net.parameters(), lr=0.01)
    train_ls, test_ls = [], []
    for _ in range(num_epochs):
        for X, y in train_iter:
            l = loss(net(X), y.view(-1, 1))
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
        train_labels = train_labels.view(-1, 1)
        test_labels = test_labels.view(-1, 1)
        train_ls.append(loss(net(train_features), train_labels).item())
        test_ls.append(loss(net(test_features), test_labels).item())

    print('final epoch: train loss', train_ls[-1], 'test loss', test_ls[-1])
    # 调用 semilogy_test 绘制训练和测试损失图像
    semilogy_test(range(1, num_epochs + 1), train_ls, 'epochs', 'loss', range(1, num_epochs + 1), test_ls,
                  ['train', 'test'])
    print('weight:', net.weight.data, '\nbias:', net.bias.data)


# 设置数据
n_train, n_test, true_w, true_b = 100, 100, [1.2, -3.4, 5.6], 5
#生成一个（200，1）的正态分布矩阵
features = torch.randn((n_train + n_test, 1))
#将三个维度的（200，1）数据拼接在一起
poly_features = torch.cat((features, torch.pow(features, 2), torch.pow(features, 3)), 1)
labels = (true_w[0] * poly_features[:, 0] + true_w[1] * poly_features[:, 1] + true_w[2] * poly_features[:, 2] + true_b)
labels += torch.tensor(np.random.normal(0, 0.01, size=labels.size()), dtype=torch.float)

num_epochs, loss = 100, torch.nn.MSELoss()

# 调用fit_and_plot进行训练和绘图
#正常情况
#fit_and_plot(poly_features[:n_train, :], poly_features[n_train:, :], labels[:n_train], labels[n_train:])
#欠拟合情况，没有进行多项式的拼接
#fit_and_plot(features[:n_train, :], features[n_train:, :], labels[:n_train],labels[n_train:])
#过拟合情况，只使用了前三行的数据
fit_and_plot(poly_features[0:2, :], poly_features[n_train:, :], labels[0:2],labels[n_train:])